#!/usr/bin/env python
"""
Simple command line tool to extract sequences of interest from a sequence file.
e.g.
python seqDB.py -db infile.sto -format stockholm -accs "acc1" "acc2" -desc "Methanobacter maripuladis" "Methanococcus"
"""
# Copyright (c) Alex Leach. 2011
# Contact: Alex Leach (albl500@york.ac.uk)
# University of York, Department of Biology,
# Wentworth Way, York YO10 4DU, United Kingdom
# 
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 3
# of the License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have receive a copy of the GNU General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>
#
# Description:
# Program to reformat or filter sequences in a variety of formats.
# See http://www.biopython.org/wiki/SeqIO for a list of formats, 
# although there are currently more supported by biopython.
# e.g. sff format (Roche 454's binary sequence format) is supported
# by biopython, but not mentioned on the webpage.

from Bio import AlignIO, SeqIO
from Bio.Align import MultipleSeqAlignment
import sys,re


options = {      '-db' : None ,
                 '-accs': [] ,
                 '-desc' : [] ,
                 '-format' : 'fasta',
                 '-out' : sys.stdout,
                 '-outformat' : 'fasta',
                 '--reverse' : False
        }

helpTxt = {     '-db' : 'File which contains sequences of interest',
                '-accs' : 'List of accessions to look for within <-db>. This is relaxed, so returns all matches which contain a given accession, even if it isnt the complete accession. Separate with spaces.',
                '-desc' : 'Descriptions to look for within db. Again, is relaxed, so returns all matches. Envelope each description with " and separate with spaces.',
                '-format' : 'The input sequence format. Default: fasta',
                '-out' : 'Output file. Default: standard out. N.B. Appends to the file.',
                '-outformat' : 'Output file format. Default: fasta.',
                '--reverse' : 'Reverse transcribe the sequences'
          }

def printHelp():
        global helpTxt
        print 'Valid options:-'
        for key,val in helpTxt.items():
                print key,' : ',val

def parseArgs(args):
        global options
        global helpTxt
        if len(args) == 0:
                printHelp()
                exit()
        prevargDash = False
        for i in xrange(len(args)):
                if args[i].startswith('--'):
                    other = options[args[i]]
                    options[args[i]] = not other
                elif args[i].startswith('-'):
                        if args[i] in options.keys():
                                prevargDash = True
                                command = args[i]
                        else:
                                print 'Invalid option: {0}'.format( args[i] )
                                printHelp()
                                exit()
                elif prevargDash:
                        if command in ['-db','-outformat','-out']:
                                options.update( { command :  args[i] } )
                                prevargDash = False
                        elif command == '-accs':
                                options[command].append(args[i])
                        elif command == '-desc':
                                options[command].append( '('+args[i] + ')')
                        elif command == '-format':
                                options.update( {command : args[i] } )
                                prevargDash = False
                else:
                        print "Don't know what to do with {0}".format( args[i] )
                        exit()
        return options
class Reader():
    def __init__(self,options):
        if len(options['-desc']) > 0:
            self.desc_reg = re.compile( '|'.join( options['-desc'] ), re.I )
        if len(options['-accs']) > 0:
            self.acc_reg = re.compile( '|'.join( options['-accs'] ), re.I )
        if options['--reverse']:
            self._yield = self.reverser
        else:
            self._yield = self.writer
        self.options = options
    def reverser(self,seq):
        seq.seq = seq.seq.reverse_complement()
        return seq

    def writer(self,seq):
        return seq
    def __iter__(self):
        Ndesc = len(options['-desc'])
        Naccs = len(options['-accs'])
        count = 0
        if Naccs == 0 and Ndesc == 0:
            printed = False
        else:
            printed = True
        for seq in AlignIO.parse( options['-db'],options['-format']  ):
            count += 1
            if Naccs > 0:
                accFind = self.acc_reg.search( seq.id )
                if accFind:
                    yield self._yield( seq )
                    printed = True
            elif Ndesc > 0:
                desc_find = self.desc_reg.search( seq.description )
                if desc_find:
                    yield self._yield( seq )
                    printed = True
            else:
                yield self._yield( seq )
        if not printed and count > 0:
                sys.stderr.write( "No matching sequences found. Here's the last sequence header as an example...\n" )
                sys.stderr.write( seq.format( options['-outformat']).splitlines()[0]  )
                sys.stderr.write('\n')
        return

def parseInput(options):
    Aligner = Reader( options )
    if options['-out'] != sys.stdout:
        out = file( options['-out'] , 'w' )
    else:
        out = sys.stdout
#    print sequences[0],type(sequences[0])
#    alignment_object = MultipleSeqAlignment( sequences )
    AlignIO.write( Aligner , out , options['-outformat'] )

if __name__ == '__main__':
        args = sys.argv[1:]
        options = parseArgs(args)
        parseInput(options)

